import numpy as np
from scipy.fftpack import fft
import matplotlib.pyplot as plt
from matplotlib.pylab import mpl

#mpl.rcParams['font.sans-serif'] = ['SimHei']  # 显示中文
#mpl.rcParams['axes.unicode_minus'] = False  # 显示负号
def fft_calc(x, f_s, x_size, nfft):
    w = np.hanning(nfft)                            # 加汉尼窗
    cnt = x_size // nfft                            # 计算数据长度可以覆盖几个窗口
    # 将输入数据长度补齐为窗口长度的整数倍，补齐数据为0
    # 此方式可以避免输入数据长度小于窗口长度
    if cnt == 0:                                    # 用0在数据尾部补齐
        x_pad = np.pad(x, (0, nfft - x_size))
    else:
        x_pad = np.pad(x, (0, x_size - cnt * nfft))
    cnt = len(x_pad) // nfft                            # 更新补齐的数据长度可以覆盖几个窗口
    # 以窗口长度计算输入数据的FFT
    tmp = []
    for i in range(cnt):                            # 窗与窗之间数据不重叠
        p = fft(w * x_pad[i * nfft:(i+1) * nfft])       # 计算加窗的FFT并乘以调整系数
        tmp.append(p)                               # 每个窗的结果
    # 将所有计算取平均值
    fft_result = np.mean(tmp, axis=0)               # 将所有窗平均得到最终结果

    # 根据采样宽度计算幅值
    amp = abs(fft_result)*2 / (nfft / 2)
    # 对直流分量额外调整
    amp[0] /= 2

    # 根据采样频率和采样点数，计算频率分辨率，并得到对应的频率坐标
    freq = np.arange(int(len(amp))) * f_s/nfft

    return amp, freq

def hlaffft_calc(x, f_s, x_size, nfft):
    w = np.hanning(nfft)                            # 加汉尼窗
    cnt = x_size // nfft                            # 计算数据长度可以覆盖几个窗口
    # 将输入数据长度补齐为窗口长度的整数倍，补齐数据为0
    # 此方式可以避免输入数据长度小于窗口长度
    if cnt == 0:                                    # 用0在数据尾部补齐
        x_pad = np.pad(x, (0, nfft - x_size))
    else:
        x_pad = np.pad(x, (0, x_size - cnt * nfft))
    cnt = len(x_pad) // nfft                            # 更新补齐的数据长度可以覆盖几个窗口
    # 以窗口长度计算输入数据的FFT
    tmp = []
    for i in range(cnt):                            # 窗与窗之间数据不重叠
        p = fft(w * x_pad[i * nfft:(i+1) * nfft])       # 计算加窗的FFT并乘以调整系数
        tmp.append(p)                               # 每个窗的结果
    # 将所有计算取平均值
    fft_result = np.mean(tmp, axis=0)               # 将所有窗平均得到最终结果

    # 根据采样宽度计算幅值
    amp = abs(fft_result)*2 / (nfft / 2)
    # 对直流分量额外调整
    amp[0] /= 2
    # 根据FFT特性，取一半频谱即可
    amp_half = amp[:int(len(amp) / 2)+1]
    # 根据采样频率和采样点数，计算频率分辨率，并得到对应的频率坐标
    freq_half = np.arange(int(len(amp) / 2)+1) * f_s/nfft

    return amp_half, freq_half

def loadtxtrarr(filename_data):
    dataset = []
    with open(filename_data) as txtData:
        lines = txtData.readlines()
        for line in lines:
            print("line=", line)
            print("len(line)=", len(line))
            linedata = line.strip().split()
            dataset.append(linedata)
        #dataset = np.array(dataset)
    return dataset

# 采样点选择1400个，因为设置的信号频率分量最高为600赫兹，根据采样定理知采样频率要大于信号频率2倍，
# 所以这里设置采样频率为1400赫兹（即一秒内有1400个采样点，一样意思的）
N = 1024
fs = N
x = np.linspace(0, 1, N)

f1 = 100   	# 信号的特征频率1
f2 = 300   # 信号的特征频率2
f3 = 500   # 信号的特征频率3
# 设置需要采样的信号，频率分量有0，200，400和600
y = 10 + 7 * np.sin(2 * np.pi * f1 * x) + 5 * np.sin(2 * np.pi * f2 * x) + 2 * np.sin(2 * np.pi * f3 * x)
#for i in range(N):
#    print("x[%d]=%f" %(i, x[i]))

s = "10.000000,19.070230,13.580103,13.956034,17.979506,12.690144,0.124638,6.647635"
l = s.split(",")

fp = open("ricker.txt", "r")
data = fp.read()
#print("data=", data)
print("len(data)=", len(data))
print("type(data)=", type(data))

fp.close()
l = data.split(",")
#print("l=", l)
print("len(l)=", len(l))
print("type(l)=", type(l))
print("l[0]=", l[0])
data_f = []
for i in range(1024):
    data_f += [float(l[i])]
#print("data_f=", data_f)
print("len(data_f)=", len(data_f))
print("type(data_f)=", type(data_f))

plt.figure(figsize=(24, 24))

'''
plt.subplot(411)
print("len(x)=", len(x))
print("type(x)=", type(x))

x = np.arange(N)  # 频率个数
print("len(x)=", len(x))
print("type(x)=", type(x))

plt.plot(x, data_f)
plt.title('input data_f')
ax = plt.gca()
plt.grid(which='both', axis='both')
'''
fft_y = fft(data_f)  # 快速傅里叶变换
x = np.arange(fs)  # 频率个数
half_x = x[range(int(N / 2))]   # 取一半区间

angle_y = np.angle(fft_y)       # 取复数的角度

abs_y = np.abs(fft_y)               # 取复数的绝对值，即复数的模(双边频谱)
normalization_y = abs_y / (N / 2)   # 归一化处理（双边频谱）
normalization_y[0] /= 2             # 归一化处理（双边频谱）
normalization_half_y = normalization_y[range(int(N / 2))]  # 由于对称性，只取一半区间（单边频谱）
x = np.arange(fs)  # 频率个数
plt.subplot(411)
plt.plot(x, abs_y, color="blue")
plt.title('py fft abs_y result')
ax = plt.gca()
plt.grid(which='both', axis='both')



fp = open("ricker_dft_amplitude.txt", "r")
data = fp.read()
#print("data=", data)
print("len(data)=", len(data))
print("type(data)=", type(data))

fp.close()
l = data.split(",")
#print("l=", l)
print("len(l)=", len(l))
print("type(l)=", type(l))
print("l[0]=", l[0])
print("l[1]=", l[1])
print("l[1024]=", l[1024])
#print("l[1025]=", l[1025])
output_data_f = []
for i in range(1024):
    output_data_f += [float(l[i])]
for i in range(1024):
    if output_data_f[i] >= output_data_f[0]:
        print(" output_data_f[%d]=%f " %(i, output_data_f[i]), end="")
#print("output_data_f=", output_data_f)
print("\nlen(data_f)=", len(output_data_f))
print("type(data_f)=", type(output_data_f))
plt.subplot(412)
x = np.arange(fs)  # 频率个数

plt.plot(x, output_data_f)
plt.title('ricker_fft_amplitude fft')
ax = plt.gca()
plt.grid(which='both', axis='both')


fp = open("out_fft.txt", "r")
data = fp.read()
#print("data=", data)
print("len(data)=", len(data))
print("type(data)=", type(data))

fp.close()
l = data.split(",")
#print("l=", l)
print("len(l)=", len(l))
print("type(l)=", type(l))
print("l[0]=", l[0])
print("l[1]=", l[1])
print("l[1024]=", l[1024])
#print("l[1025]=", l[1025])
output_data_f = []
for i in range(1024):
    output_data_f += [float(l[i])]
for i in range(1024):
    if output_data_f[i] >= output_data_f[0]:
        print(" output_data_f[%d]=%f " %(i, output_data_f[i]), end="")
#print("output_data_f=", output_data_f)
print("\nlen(data_f)=", len(output_data_f))
print("type(data_f)=", type(output_data_f))
print("out_fft output_data_f[1023]=", output_data_f[1023])

plt.subplot(413)
x = np.arange(fs)  # 频率个数

plt.plot(x, output_data_f)
plt.title('out_fft fft')
ax = plt.gca()
plt.grid(which='both', axis='both')

fp = open("out_fftr.txt", "r")
data = fp.read()
#print("data=", data)
print("len(data)=", len(data))
print("type(data)=", type(data))

fp.close()
l = data.split(",")
#print("l=", l)
print("len(l)=", len(l))
print("type(l)=", type(l))
print("l[0]=", l[0])
print("l[1]=", l[1])
print("l[1024]=", l[1024])
#print("l[1025]=", l[1025])
output_data_f = []
for i in range(1024):
    output_data_f += [float(l[i])]
for i in range(1024):
    if output_data_f[i] >= output_data_f[0]:
        print(" output_data_f[%d]=%f " %(i, output_data_f[i]), end="")
#print("output_data_f=", output_data_f)
print("\nlen(data_f)=", len(output_data_f))
print("type(data_f)=", type(output_data_f))
print("out_fft output_data_f[1023]=", output_data_f[1023])

plt.subplot(414)
x = np.arange(fs)  # 频率个数

plt.plot(x, output_data_f)
plt.title('out_fftr fft')
ax = plt.gca()
plt.grid(which='both', axis='both')


plt.show()
exit()




fp = open("cpp1.txt", "r")
data = fp.read()
#print("data=", data)
print("len(data)=", len(data))
print("type(data)=", type(data))

fp.close()
l = data.split(",")
#print("l=", l)
print("len(l)=", len(l))
print("type(l)=", type(l))
print("l[0]=", l[0])
print("l[1]=", l[1])
print("l[1024]=", l[1024])
#print("l[1025]=", l[1025])
fft_output_data_f = []
for i in range(1024):
    fft_output_data_f += [float(l[i])]
for i in range(1024):
    if fft_output_data_f[i] >= fft_output_data_f[0]:
        print(" fft_output_data_f[%d]=%f " %(i, fft_output_data_f[i]), end="")
print("fft_output_data_f=", fft_output_data_f)
print("\nlen(fft_output_data_f)=", len(fft_output_data_f))
print("type(fft_output_data_f)=", type(fft_output_data_f))
plt.subplot(412)
x = np.arange(fs)  # 频率个数
plt.plot(x, fft_output_data_f)
plt.title('cpp1 fft_output_data_f')
ax = plt.gca()
plt.grid(which='both', axis='both')


fp = open("cpp2.txt", "r")
data = fp.read()
#print("data=", data)
print("len(data)=", len(data))
print("type(data)=", type(data))

fp.close()
l = data.split(",")
#print("l=", l)
print("len(l)=", len(l))
print("type(l)=", type(l))
print("l[0]=", l[0])
print("l[1]=", l[1])
print("l[1024]=", l[1024])
#print("l[1025]=", l[1025])
fft_output_data_f = []
for i in range(1024):
    fft_output_data_f += [float(l[i])]
for i in range(1024):
    if fft_output_data_f[i] >= fft_output_data_f[0]:
        print(" fft_output_data_f[%d]=%f " %(i, fft_output_data_f[i]), end="")
print("fft_output_data_f=", fft_output_data_f)
print("\nlen(fft_output_data_f)=", len(fft_output_data_f))
print("type(fft_output_data_f)=", type(fft_output_data_f))
plt.subplot(413)
x = np.arange(fs)  # 频率个数
plt.plot(x, fft_output_data_f)
plt.title('cpp2 fft_output_data_f')
ax = plt.gca()
plt.grid(which='both', axis='both')



'''
for i in range(1024):
    print(" output_data_f[%d]=%f " %(i, output_data_f[i]), end="")
    print(" fft_output_data_f[%d]=%f \n" %(i, fft_output_data_f[i]), end="")
'''


fft_y = fft(data_f)  # 快速傅里叶变换
x = np.arange(fs)  # 频率个数
half_x = x[range(int(N / 2))]   # 取一半区间

angle_y = np.angle(fft_y)       # 取复数的角度

abs_y = np.abs(fft_y)               # 取复数的绝对值，即复数的模(双边频谱)
normalization_y = abs_y / (N / 2)   # 归一化处理（双边频谱）
normalization_y[0] /= 2             # 归一化处理（双边频谱）
normalization_half_y = normalization_y[range(int(N / 2))]  # 由于对称性，只取一半区间（单边频谱）
x = np.arange(fs)  # 频率个数
plt.subplot(414)
plt.plot(x, normalization_y, color="blue")
plt.title('normalization_y result')
ax = plt.gca()
plt.grid(which='both', axis='both')
plt.show()

exit()
nfft = fs                                                        				# FFT的窗长
(amp, freq) = fft_calc(data_f, fs , N , nfft)
print("amp[0]=", amp[0])
print("amp[1023]=", amp[1023])
print("amp=", amp)
for i in range(1024):
    if amp[i] >= amp[0]:
        print(" amp[%d]=%f " %(i, amp[i]), end="")

print("\nfft_calc len(amp)=", len(amp))
print("fft_calc type(amp)=", type(amp))
# plot for FFT
#fig = plt.figure(figsize=(12, 8))
plt.subplot(414)
plt.plot(freq, amp, color="blue")

plt.title('fft_calc-1 linspace result')
plt.xlabel('Freq /Hz')
plt.ylabel('urad/Hz')

ax = plt.gca()
plt.grid(which='both', axis='both')



fft_y = fft(y)  # 快速傅里叶变换

x = np.arange(fs)  # 频率个数
half_x = x[range(int(N / 2))]   # 取一半区间

angle_y = np.angle(fft_y)       # 取复数的角度

abs_y = np.abs(fft_y)               # 取复数的绝对值，即复数的模(双边频谱)
normalization_y = abs_y / (N / 2)   # 归一化处理（双边频谱）
normalization_y[0] /= 2             # 归一化处理（双边频谱）
normalization_half_y = normalization_y[range(int(N / 2))]  # 由于对称性，只取一半区间（单边频谱）

for i in range(1024):
    #if abs_y[i] >= abs_y[0]:
    print("abs_y[%d]=%f" %(i, abs_y[i]), end="")
# 绘制结果
plt.figure()

plt.subplot(231)
plt.plot(x, data_f)
plt.title('原始波形')

plt.subplot(232)
plt.plot(x, fft_y, 'black')
plt.title('双边振幅谱(未求振幅绝对值)', fontsize=9, color='black')

plt.subplot(211)
plt.plot(x, abs_y, 'r')
print("abs_y=", abs_y)
print("len(abs_y)=", len(abs_y))
print("type(abs_y)=", type(abs_y))
plt.title('双边振幅谱(未归一化)', fontsize=9, color='red')

plt.subplot(212)
#plt.plot(x, angle_y, 'violet')
plt.plot(x, output_data_f, 'yellow')

plt.title('双边相位谱(未归一化)', fontsize=9, color='violet')


plt.subplot(235)
plt.plot(x, normalization_y, 'g')

plt.title('双边振幅谱(归一化)', fontsize=9, color='green')

plt.subplot(236)
plt.plot(half_x, normalization_half_y, 'blue')

plt.title('单边振幅谱(归一化)', fontsize=9, color='blue')

plt.show()
